# coding=utf-8
from __future__ import print_function, absolute_import
from gm.api import *
import pandas as pd
import time
import datetime
import tushare as ts
from concurrent.futures import *
import tqdm
import chage_data_type
import multiprocessing as mp
import os


class GetData(object):
    def __init__(self):
        super(GetData, self).__init__()
        set_token('df082f274b3ee380fcd2a98a2e5a7217dac1268b')
        self.pro = ts.pro_api('762be0e261faf0fc1d7807e160da4506b2268d600ce6e0da581cc1b9')
        down_day = 600
        i = get_instruments(symbols=None, exchanges=None, sec_types='1', names=None, skip_suspended=False,
                            skip_st=False,
                            fields=None, df=True)
        i = i[i['trade_date'] - i['listed_date'] >= datetime.timedelta(down_day)]
        i = i[~(i['symbol'].str.startswith('SHSE.900') ^ i['symbol'].str.startswith('SZSE.200') ^
                i['symbol'].str.startswith('SHSE.688') ^ i['symbol'].str.startswith('SZSE.300'))]
        self.code_ts = i['symbol']
        i = i.set_index('symbol').sort_values('listed_date', ascending=False)
        self.codes = list(i.index)
        d = history_n('SZSE.399300', '1d', down_day, end_time=None, fields=None, skip_suspended=True, fill_missing=None,
                      adjust=ADJUST_NONE, adjust_end_time='', df=True)['bob'].dt.strftime('%Y-%m-%d')
        self.d_gd = d.tolist()  # .dt.strftime('%Y-%m-%d')
        self.d_ts = list(d.astype('datetime64[ns]').dt.strftime('%Y%m%d'))

    def daily_ts(self):
        def get_daily_ts(date):
            time.sleep(1)
            df = self.pro.daily(trade_date=date)
            df = df[~(df['ts_code'].str.startswith('900') ^ df['ts_code'].str.startswith('200') ^
                      df['ts_code'].str.startswith('688') ^ df['ts_code'].str.startswith('300'))]
            df = df.rename(columns={'ts_code': 'symbol'})
            df['symbol'] = df['symbol'].apply(self.change_day)
            df['symbol'] = df['symbol'].map(
                lambda code: 'SZSE.' + code if code.startswith('3') or code.startswith('0') else 'SHSE.' + code)
            df['trade_date'] = df['trade_date'].astype('datetime64[ns]')
            df['vol'] = df['vol'].astype('uint64')
            df['amount'] = df['amount'].astype('uint64')
            df = df[['symbol', 'trade_date', 'open', 'high', 'low', 'close', 'pre_close', 'vol', 'amount']]
            df = df.dropna()
            return df

        num = 450
        try:
            print('开始读取本地daily数据库...')
            downed = pd.read_hdf('D:/data/history_daily.h5', 'history_daily', columns=['trade_date']
                                 ).trade_date.dt.strftime('%Y-%m-%d').drop_duplicates().tolist()
            dif = list(set(self.d_gd).difference(set(downed)))
            dif = pd.Series(dif).astype('datetime64[ns]').dt.strftime('%Y%m%d').tolist()
            if len(dif) >= num:
                print('一共 ' + str(len(dif)) + ' 个daily数据需要更新...')
                dif = [dif[n:n + num] for n in range(0, len(dif), num)]
                print('开始分 ' + str(len(dif)) + ' 批下载...')
                for x in dif:
                    print('正在下载第 ' + str(dif.index(x) + 1) + ' 批数据,请稍候...')
                    all_task = [ThreadPoolExecutor(4).submit(get_daily_ts, date) for date in x]
                    with tqdm.tqdm(total=(len(x))) as tq:
                        for future in as_completed(all_task):
                            data = future.result()
                            data.to_hdf('D:/data/history_daily.h5', 'history_daily', append=True, format='t',
                                        index=False,
                                        data_columns=['symbol', 'trade_date'])
                            tq.set_description('history ')
                            tq.update()
                    print('准备开始下一阶段的下载任务!')
                    if len(dif) - dif.index(x) > 1:
                        time.sleep(60)
                print('daily数据更新完毕!!!')
            elif num > len(dif) > 0:
                dif = dif
                print('一共 ' + str(len(dif)) + ' 个daliy数据需要更新...')
                all_task = [ThreadPoolExecutor(4).submit(get_daily_ts, date) for date in dif]
                with tqdm.tqdm(total=(len(dif))) as tq:
                    for future in as_completed(all_task):
                        data = future.result()
                        data.to_hdf('D:/data/history_daily.h5', 'history_daily', append=True, format='t', index=False,
                                    data_columns=['symbol', 'trade_date'])
                        tq.set_description('history ')
                        tq.update()
                print('daily数据更新完毕!!!')
            else:
                print('daily没有数据需要更新!!!')
        except FileNotFoundError:
            dif = self.d_ts
            if len(dif) >= num:
                print('一共 ' + str(len(dif)) + ' 个daily数据需要下载...')
                dif = [dif[n:n + num] for n in range(0, len(dif), num)]
                print('开始分 ' + str(len(dif)) + ' 批下载...')
                for x in dif:
                    print('正在下载第 ' + str(dif.index(x) + 1) + ' 批数据,请稍候...')
                    all_task = [ThreadPoolExecutor(4).submit(get_daily_ts, date) for date in x]
                    with tqdm.tqdm(total=(len(x))) as tq:
                        for future in as_completed(all_task):
                            data = future.result()
                            data.to_hdf('D:/data/history_daily.h5', 'history_daily', append=True, format='t',
                                        index=False,
                                        data_columns=['symbol', 'trade_date'])
                            tq.set_description('history ')
                            tq.update()
                    print('准备开始下一阶段的下载任务!')
                    if len(dif) - dif.index(x) > 1:
                        time.sleep(60)
                print('daily数据下载完毕!!!')
            elif num > len(dif) > 0:
                dif = dif
                print('一共 ' + str(len(dif)) + ' 个daliy数据需要更新...')
                all_task = [ThreadPoolExecutor(4).submit(get_daily_ts, date) for date in dif]
                with tqdm.tqdm(total=(len(dif))) as tq:
                    for future in as_completed(all_task):
                        data = future.result()
                        data.to_hdf('D:/data/history_daily.h5', 'history_daily', append=True, format='t', index=False,
                                    data_columns=['symbol', 'trade_date'])
                        tq.set_description('history ')
                        tq.update()
                print('daily数据下载完毕!!!')

    def change_day(self, code):
        code_split = code.split('.', 1)
        return code_split[0]

    def instruments(self):
        def get_ih(x):
            time.sleep(0.5)
            ih = get_history_instruments(self.codes, fields=None, start_date=x, end_date=x, df=True)
            ih = ih[ih['is_suspended'] == 0]
            ih = ih[['symbol', 'trade_date', 'upper_limit', 'lower_limit', 'sec_level', 'is_suspended', 'pre_close']]
            ih = ih.dropna()
            return ih

        def tpe(dif):
            with ThreadPoolExecutor(8) as etr:
                all_task = [etr.submit(get_ih, x) for x in dif]
                with tqdm.tqdm(total=(len(dif))) as tq:
                    for future in as_completed(all_task):
                        data = future.result()
                        data['trade_date'] = data['trade_date'].dt.date.astype('datetime64[ns]')
                        data.to_hdf('D:/data/history_instruments.h5', 'history_instruments', append=True, format='t',
                                    index=False,
                                    data_columns=['symbol', 'trade_date'])
                        tq.set_description('instruments ')
                        tq.update()

        try:
            print('开始读取本地instruments数据库...')
            downed = pd.read_hdf('D:/data/history_instruments.h5', 'history_instruments',
                                 columns=['trade_date']).trade_date.dt.strftime('%Y-%m-%d').drop_duplicates().tolist()
            dif = list(set(self.d_gd).difference(set(downed)))
            if not dif:
                print('instruments没有数据需要更新!!')
            else:
                print('一共 ' + str(len(dif)) + ' 个instruments数据需要更新...')
                tpe(dif)
            '''
            for x in dif:
                data = get_ih(x)
                data['trade_date'] = data['trade_date'].dt.date.astype(
                    'datetime64[ns]')  # .astype('datetime64[ns]')
                data.to_hdf('D:/data/history_instruments.h5', 'history_instruments', append=True, format='t', index=False,
                            data_columns=['symbol', 'trade_date'])
            '''
        except FileNotFoundError:
            dif = self.d_gd
            print('一共 ' + str(len(dif)) + ' 个instruments数据需要下载...')
            tpe(dif)
            print('instruments数据下载完毕!')

    def turn_gm(self):
        print('turn数据使用掘金数据源,稳定...')

        def get_turn_gm(code):
            time.sleep(1)
            turn = get_fundamentals(table='trading_derivative_indicator', symbols=code, start_date=date,
                                    end_date=end,
                                    fields='TURNRATE', filter=None, order_by=None, df=True)
            # turn = get_fundamentals_n(table='trading_derivative_indicator', symbols=codes, end_date=end, count=down_day,
            # fields='TURNRATE', filter=None, order_by=None, df=True)
            turn = turn.dropna()
            return turn

        def tpe(codes):
            with tqdm.tqdm(total=len(codes)) as tq:
                with ThreadPoolExecutor(10) as etr:
                    all_task = [etr.submit(get_turn_gm, code) for code in codes]
                    for future in as_completed(all_task):
                        data = future.result()
                        try:
                            data['trade_date'] = data['pub_date'].dt.strftime('%Y-%m-%d').astype('datetime64[ns]')
                            data = data[['symbol', 'trade_date', 'TURNRATE']]
                            data.to_hdf('D:/data/turn_daily_gm.h5', 'turn_daily_gm', append=True, format='t',
                                        index=False,
                                        data_columns=['symbol', 'trade_date'])
                        except KeyError:
                            print('bug!')
                        tq.set_description('turn_gm ')
                        tq.update()

        try:
            print('开始读取本地turnover数据库...')
            data = pd.read_hdf('D:/data/turn_daily_gm.h5', 'turn_daily_gm')
            downed = data.trade_date.dt.strftime('%Y-%m-%d').drop_duplicates().tolist()
            downed_symbol = data.symbol.drop_duplicates().tolist()
            dif = list(set(self.d_gd).difference(set(downed)))
            undown_symbol = list(set(self.codes).difference(set(downed_symbol)))
            if dif:
                print('一共 ' + str(len(dif)) + ' 个交易日turnover数据需要更新...')
                date = dif[-1]
                end = self.d_gd[-1]
                tpe(self.codes)
            if undown_symbol:
                print('一共 ' + str(len(undown_symbol)) + ' 个股票turnover数据需要更新...')
                date = self.d_gd[0]
                end = self.d_gd[-1]

                '''
                for x in codes:
                    print('开始下载')
                    data = get_turn_gm(x)
                    data['trade_date'] = data['pub_date'].dt.strftime('%Y-%m-%d').astype('datetime64[ns]')
                    data = data[['symbol', 'trade_date', 'TURNRATE']]
                    data.to_hdf('D:/data/turn_daily_gm.h5', 'turn_daily_gm', append=True, format='t', index=False,
                                data_columns=['symbol', 'trade_date'])
                '''
                tpe(undown_symbol)
                print('turnover数据更新完毕!!!')
            elif not dif or not undown_symbol:
                print('turnover没有数据需要更新')
        except FileNotFoundError:
            date = self.d_gd[0]
            end = self.d_gd[-1]
            print('一共 ' + str(len(self.codes)) + ' 个turnover数据需要下载...')
            print()
            tpe(self.codes)
            '''
            for x in codes:
                data = get_turn_gm(x)
                data['trade_date'] = data['pub_date'].dt.strftime('%Y-%m-%d').astype('datetime64[ns]')
                data = data[['symbol', 'trade_date', 'TURNRATE']]
                data.to_hdf('D:/data/turn_daily_gm.h5', 'turn_daily_gm', append=True, format='t', index=False,
                            data_columns=['symbol', 'trade_date'])
            '''
            print('turnover数据下载完毕!!!')


if __name__ == '__main__':
    getdata = GetData()

